from flask import Flask, jsonify, request
from agent.hub_agent import agent_tui, bot

app = Flask(__name__)

@app.route('/api/ai-hub', methods=['GET', 'POST'])
def ai_hub():
    if request.method == 'POST':
        msg = request.json.get('msg')
        pic = request.json.get('image')
        if pic is not None and not pic.startswith('data:image/'):
            return jsonify({'message': 'image format error, should start with data:image/base64 and base64encode'})
        post_msg = []
        if pic:
            post_msg.append({'image': pic})
        post_msg.append({'text': msg})

        response = agent_tui(bot, post_msg)
        res_msg = 'no return msg by llm'
        for res in response:
            # 判断 res 中是否有 function_call 字段
            if 'function_call' in res and res['function_call']:
                continue
            if res['role'] == 'assistant' and res['name'] == 'AiHubBot':
                res_msg = res['content']
            else:
                continue
        return jsonify({'message': res_msg})
    return jsonify({'message': 'AI Hub API is working'})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5001, debug=True)